from flask import Flask, render_template, request, jsonify, send_from_directory
from flask_cors import CORS
from flask_login import LoginManager, current_user, login_user, logout_user, login_required
from werkzeug.utils import secure_filename
import os
from dotenv import load_dotenv
from models.database import db, User, Image, Prediction

# 使用条件导入
try:
    from routes.auth import auth_bp
    auth_available = True
except ImportError as e:
    print(f"警告: 无法导入auth模块: {e}")
    auth_available = False

try:
    from routes.data import data_bp
    data_available = True
except ImportError as e:
    print(f"警告: 无法导入data模块: {e}")
    data_available = False

try:
    from routes.prediction import prediction_bp
    prediction_available = True
except ImportError as e:
    print(f"警告: 无法导入prediction模块: {e}")
    prediction_available = False

try:
    from routes.analysis import analysis_bp
    analysis_available = True
except ImportError as e:
    print(f"警告: 无法导入analysis模块: {e}")
    analysis_available = False

load_dotenv()

app = Flask(__name__)
app.config['SECRET_KEY'] = os.getenv('SECRET_KEY', 'default-secret-key')
app.config['SQLALCHEMY_DATABASE_URI'] = os.getenv('DATABASE_URI', 'sqlite:///river_segmentation.db')
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False
app.config['UPLOAD_FOLDER'] = os.path.join(os.getcwd(), 'uploads')
app.config['MODEL_FOLDER'] = os.path.join(os.getcwd(), 'models', 'weights')
app.config['RESULTS_FOLDER'] = os.path.join(os.getcwd(), 'results')

# 确保必要的文件夹存在
for folder in [app.config['UPLOAD_FOLDER'], app.config['MODEL_FOLDER'], app.config['RESULTS_FOLDER']]:
    os.makedirs(folder, exist_ok=True)

# 初始化数据库
db.init_app(app)

# 设置登录管理器
login_manager = LoginManager()
login_manager.init_app(app)
login_manager.login_view = 'auth.login'

@login_manager.user_loader
def load_user(user_id):
    return User.query.get(int(user_id))

# 条件注册蓝图
if auth_available:
    app.register_blueprint(auth_bp)
if data_available:
    app.register_blueprint(data_bp)
if prediction_available:
    app.register_blueprint(prediction_bp)
if analysis_available:
    app.register_blueprint(analysis_bp)

# 允许跨域请求
CORS(app)

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/dashboard')
@login_required
def dashboard():
    # 从数据库获取实际统计数据
    # 获取当前用户上传的图像数量
    image_count = Image.query.filter_by(user_id=current_user.id).count()
    
    # 获取用户所有图像的ID
    user_images = Image.query.filter_by(user_id=current_user.id).all()
    image_ids = [img.id for img in user_images]
    
    # 获取预测结果数量
    prediction_count = 0
    max_water_percentage = 0
    
    if image_ids:
        # 查询与用户图像相关的所有预测结果
        predictions = Prediction.query.filter(Prediction.image_id.in_(image_ids)).all()
        prediction_count = len(predictions)
        
        # 计算最大水域占比
        for pred in predictions:
            if pred.water_area and pred.water_area > max_water_percentage:
                max_water_percentage = pred.water_area
    
    # 创建统计数据
    stats = {
        'image_count': image_count,
        'prediction_count': prediction_count,
        'max_water_percentage': round(max_water_percentage * 100)  # 转换为百分比
    }
    
    # 获取最近活动（最近5条预测记录）
    recent_activities = []
    if image_ids:
        recent_predictions = Prediction.query.filter(
            Prediction.image_id.in_(image_ids)
        ).order_by(Prediction.prediction_date.desc()).limit(5).all()
        
        for pred in recent_predictions:
            image = Image.query.get(pred.image_id)
            if image:
                activity = {
                    'type': '预测',
                    'model': pred.model_name,
                    'image_name': image.filename.split('_', 1)[1] if '_' in image.filename else image.filename,
                    'date': pred.prediction_date.strftime('%Y-%m-%d %H:%M'),
                    'prediction_id': pred.id,
                    'image_id': image.id
                }
                recent_activities.append(activity)
    
    return render_template('dashboard.html', stats=stats, recent_activities=recent_activities)

@app.route('/uploads/<filename>')
@login_required
def uploaded_file(filename):
    return send_from_directory(app.config['UPLOAD_FOLDER'], filename)

@app.route('/static/results/<filename>')
@login_required
def result_file(filename):
    return send_from_directory(app.config['RESULTS_FOLDER'], filename)

if __name__ == '__main__':
    with app.app_context():
        db.create_all()
    app.run(debug=True) 